import numpy as np
from typing_extensions import Iterator
from ...base import DatabaseBuilder
import tqdm


def latlon2pix(lat, lon, radius, depth):
    """
    将经纬度转换为球面坐标，注意单位是弧度
    """
    x = (radius + depth) * np.cos(lat) * np.cos(lon)
    y = (radius + depth) * np.cos(lat) * np.sin(lon)
    z = (radius + depth) * np.sin(lat)
    return x, y, z


class TriadHierarchyBuilder(DatabaseBuilder):

    def build(self, data: Iterator, max_field=None, hierarch_r=None) -> np.ndarray:
        """
        基于层次法的识别库构建
        """
        max_field = np.inf if max_field is None else max_field
        # 计算真实三维点中心
        x, y, z = latlon2pix(
            data[0] * np.pi / 180, data[1] * np.pi / 180, self.radius, data[3]
        )
        r = data[2]
        # 计算着陆平面
        plane = np.array((x, y, z, np.ones_like(x))).T
        _, _, Vh = np.linalg.svd(plane)
        # 平面方程
        plane = Vh[-1]
        ################
        # 数值不稳定，改为直接计算
        # # 空间中以球面表达的陨石坑矩阵表示, N * 4 * 4
        # Q = np.array(
        #     [
        #         [np.ones_like(x), np.zeros_like(x), np.zeros_like(x), -x],
        #         [np.zeros_like(x), np.ones_like(x), np.zeros_like(x), -y],
        #         [np.zeros_like(x), np.zeros_like(x), np.ones_like(x), -z],
        #         [-x, -y, -z, x * x + y * y + z * z - r * r],
        #     ]
        # ).T
        # # 计算平面上的陨石坑矩阵表示，N * 3 * 3
        # C = cone_plane_cross(Q, plane)
        # 计算平面上的陨石坑的中心
        # center = center_ellipse(C)
        # center = np.array(center).T
        ################
        # 将中心的空间坐标投影为平面坐标：
        # 构建旋转矩阵，使得法向量对齐到z轴
        u = plane[:3] / np.linalg.norm(plane[:3])
        v = np.cross(u, np.array([1, 0, 0]))
        w = np.cross(u, v)
        w = w / np.linalg.norm(w)
        v = np.cross(w, u)
        R = np.column_stack((v, w, u))
        # 旋转中心
        T = -R @ np.array([0, 0, 1])
        # 旋转矩阵
        R = R.T
        # 旋转后的中心
        C = R @ np.array((x, y, z)) + T[:, None]
        center = C[:2].T
        # 计算直径的前30%分位数和后30%分位数
        if hierarch_r is None:
            hierarch_r = np.percentile(r, 30), np.percentile(r, 70)
        ###############
        C = np.array(
            [
                [np.ones_like(C[0]), np.zeros_like(C[0]), -C[0]],
                [np.zeros_like(C[0]), np.ones_like(C[0]), -C[1]],
                [-C[0], -C[1], C[0] * C[0] + C[1] * C[1] - r * r],
            ]
        ).T
        # 分层的不变量
        invariants = []
        # 逐层计算不变量
        index = np.arange(len(center))
        for i in tqdm.tqdm(index):
            # 计算当前中心的视场范围，仅取序号大于当前值的，保证不重复计算 TODO
            idx_i = np.argwhere(
                (np.linalg.norm(center - center[i], axis=1) < max_field) & (index > i)
            )
            if len(idx_i) < 2:
                continue
            idx_i = idx_i.squeeze()
            for it_j, j in enumerate(idx_i):
                # 并行计算不变量，每个描述子的存储方式为按角度从大到小依次存储
                k = idx_i[it_j + 1 :]
                if len(k) == 0:
                    continue
                ix, Is = self.invariant(i, j, k, C[i], C[j], C[k])
                invariants.append((*ix, *Is))
        invariants = np.concatenate(invariants, axis=1)
        # 清除位于1-3层次的全部三元组
        idx = invariants[:3].astype(np.int32)
        remain_ind = ~(
            (  # 1-1-3
                (idx[0] < hierarch_r[0])
                & (idx[1] < hierarch_r[0])
                & (idx[2] > hierarch_r[1])
            )
            | (  # 1-3-1
                (idx[0] < hierarch_r[0])
                & (idx[1] > hierarch_r[1])
                & (idx[2] < hierarch_r[0])
            )
            | (  # 3-1-1
                (idx[0] > hierarch_r[1])
                & (idx[1] < hierarch_r[0])
                & (idx[2] < hierarch_r[0])
            )
            | (  # 1-3-3
                (idx[0] < hierarch_r[0])
                & (idx[1] > hierarch_r[1])
                & (idx[2] > hierarch_r[1])
            )
            | (  # 3-3-1
                (idx[0] > hierarch_r[1])
                & (idx[1] > hierarch_r[1])
                & (idx[2] < hierarch_r[0])
            )
            | (  # 3-1-3
                (idx[0] > hierarch_r[1])
                & (idx[1] < hierarch_r[0])
                & (idx[2] > hierarch_r[1])
            )
        )
        return {
            "catalog": np.block([[x], [y], [z]]).T,
            "temp": np.block([center, r[:, None]]),
            "idx": idx[:, remain_ind],
            "descriptor": invariants[3:, remain_ind],
            "constants": {"max_field": max_field, "hierarch": hierarch_r},
        }

    def invariant(self, i, j, k, Q1, Q2, Q3):
        raise NotImplementedError
